package biz

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"
	"transfer-service/internal/global"

	"github.com/gin-gonic/gin"
	"github.com/go-kratos/kratos/v2/log"
	"github.com/go-kratos/kratos/v2/transport/http"
	"github.com/gorilla/websocket"
)

// Client represents a WebSocket client connection
type Client struct {
	conn *websocket.Conn
	send chan []byte
}

// Message represents the message structure exchanged between client and server
type Message struct {
	Type string `json:"type"`
	Data any    `json:"data"`
}

const (
	// Connection timeouts
	heartbeatInterval = 5 * time.Second
	writeWait         = 3 * time.Second
	readWait          = 60 * time.Second
	maxMessageSize    = 4096
)

// WebsocketUseCase manages WebSocket connections
type WebsocketUseCase struct {
	log     *log.Helper
	clients map[*Client]bool
	mu      sync.RWMutex
	done    chan struct{}
	// Removed pub/sub functionality
}

// NewWebsocketUseCase creates a new WebSocket use case
func NewWebsocketUseCase(logger log.Logger) *WebsocketUseCase {
	uc := &WebsocketUseCase{
		log:     log.NewHelper(log.With(logger, "module", "biz/Websocket")),
		clients: make(map[*Client]bool),
		done:    make(chan struct{}),
	}
	return uc
}

// WsHandler handles WebSocket connection requests
func (uc *WebsocketUseCase) WsHandler(c *gin.Context) {
	upgrader := websocket.Upgrader{
		ReadBufferSize:  maxMessageSize,
		WriteBufferSize: maxMessageSize,
		CheckOrigin:     func(r *http.Request) bool { return true },
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		uc.log.Error("Failed to upgrade connection:", err)
		return
	}

	client := &Client{
		conn: conn,
		send: make(chan []byte, 256),
	}

	// Register new client
	uc.mu.Lock()
	uc.clients[client] = true
	uc.mu.Unlock()

	uc.log.Info("New client connected")

	// Start client goroutines
	go uc.writePump(client)
	go uc.readPump(client)
}

// unregisterClient removes a client and closes connection
func (uc *WebsocketUseCase) unregisterClient(client *Client) {
	uc.mu.Lock()
	defer uc.mu.Unlock()

	if _, ok := uc.clients[client]; ok {
		delete(uc.clients, client)
		close(client.send)
		err := client.conn.Close()
		if err != nil {
			global.LOG.Error("Failed to close client:", err)
			return
		}
		uc.log.Info("Client disconnected")
	}
}

// writePump handles sending messages to the client
func (uc *WebsocketUseCase) writePump(client *Client) {
	ticker := time.NewTicker(heartbeatInterval)
	defer func() {
		ticker.Stop()
		uc.unregisterClient(client)
	}()

	for {
		select {
		case message, ok := <-client.send:
			if !ok {
				// Channel closed, connection is closing
				err := client.conn.WriteMessage(websocket.CloseMessage, []byte{})
				if err != nil {
					global.LOG.Error("Failed to close client:", err)
					return
				}
				return
			}

			err := client.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err != nil {
				global.LOG.Error("Failed to set write deadline:", err)
				return
			}
			w, err := client.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				uc.log.Error("Failed to get writer:", err)
				return
			}

			_, err = w.Write(message)
			if err != nil {
				global.LOG.Error("Failed to write message:", err)
				return
			}

			// Add queued messages to the current WebSocket message
			n := len(client.send)
			for i := 0; i < n; i++ {
				_, err := w.Write(<-client.send)
				if err != nil {
					global.LOG.Error("Failed to write message:", err)
					return
				}
			}

			if err := w.Close(); err != nil {
				uc.log.Error("Failed to close writer:", err)
				return
			}

		case <-ticker.C:
			err := client.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err != nil {
				global.LOG.Error("Failed to set write deadline:", err)
				return
			}
			if err := client.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				uc.log.Error("Heartbeat failed:", err)
				return
			}
			uc.log.Debug("Sent PING")
		}
	}
}

// readPump handles receiving messages from the client
func (uc *WebsocketUseCase) readPump(client *Client) {
	defer uc.unregisterClient(client)

	client.conn.SetReadLimit(maxMessageSize)
	err := client.conn.SetReadDeadline(time.Now().Add(readWait))
	if err != nil {
		global.LOG.Error("Failed to set read deadline:", err)
		return
	}

	// Setup ping-pong handlers
	client.conn.SetPongHandler(func(string) error {
		uc.log.Debug("Received PONG")
		return client.conn.SetReadDeadline(time.Now().Add(readWait))
	})

	client.conn.SetPingHandler(func(string) error {
		uc.log.Debug("Received PING")
		err := client.conn.SetWriteDeadline(time.Now().Add(writeWait))
		if err != nil {
			global.LOG.Error("Failed to set write deadline:", err)
			return err
		}
		err = client.conn.WriteMessage(websocket.PongMessage, nil)
		if err == nil {
			uc.log.Debug("Sent PONG")
		}
		return client.conn.SetReadDeadline(time.Now().Add(readWait))
	})

	for {
		_, message, err := client.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				uc.log.Errorf("Connection closed unexpectedly: %v", err)
			}
			break
		}

		// Handle simple ping message
		if string(message) == "ping" {
			client.send <- []byte("pong")
			continue
		}

		// Process incoming message
		var msg Message
		if err := json.Unmarshal(message, &msg); err != nil {
			uc.log.Error("Failed to unmarshal message:", err)
			continue
		}

		// Process message and send response
		uc.handleMessage(client, msg)
	}
}

// handleMessage processes messages received from clients
func (uc *WebsocketUseCase) handleMessage(client *Client, msg Message) {
	uc.log.Infof("Received message of type: %s", msg.Type)

	// Create response - example implementation
	response := Message{
		Type: "response",
		Data: msg.Data,
	}

	responseData, err := json.Marshal(response)
	if err != nil {
		uc.log.Error("Failed to marshal response:", err)
		return
	}

	select {
	case client.send <- responseData:
		// Message sent to client's queue
	default:
		// Client send buffer is full, close connection
		uc.log.Error("Client send buffer full")
		uc.unregisterClient(client)
	}
}

// SendToClient sends a message to a specific client
func (uc *WebsocketUseCase) SendToClient(client *Client, msgType string, data any) error {
	message, err := json.Marshal(Message{Type: msgType, Data: data})
	if err != nil {
		return fmt.Errorf("failed to marshal message: %w", err)
	}

	select {
	case client.send <- message:
		return nil
	default:
		uc.unregisterClient(client)
		return fmt.Errorf("client send buffer full")
	}
}

// BroadcastMessage sends a message to all connected clients
func (uc *WebsocketUseCase) BroadcastMessage(data any) error {
	message, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to marshal message: %w", err)
	}

	uc.mu.RLock()
	defer uc.mu.RUnlock()

	for client := range uc.clients {
		select {
		case client.send <- message:
			// Message queued
		default:
			// Buffer full, close connection
			go uc.unregisterClient(client)
		}
	}
	return nil
}

// Close shuts down the WebSocket service
func (uc *WebsocketUseCase) Close() {
	close(uc.done)

	uc.mu.Lock()
	for client := range uc.clients {
		err := client.conn.Close()
		if err != nil {
			global.LOG.Error("Failed to close client:", err)
			return
		}
		delete(uc.clients, client)
	}
	uc.mu.Unlock()

	uc.log.Info("WebSocket service shut down")
}
